"""
RAGFlow MCP工具模块

该模块提供了三种主要工具功能：
1. 列出所有可用的助手列表
2. 使用指定的助手回答问题
3. 自动分析并选择合适的助手回答问题

通过调用RAGutils中的函数实现核心功能。
"""
import logging
import os
import time

from dotenv import load_dotenv
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.runnables import Runnable
from langchain_openai import ChatOpenAI

from RAGFlow_utils.create_ask_delete import create_ask_delete
from RAGFlow_utils.list_chat_assistant import get_assistant_list

# 获取项目根目录的绝对路径
PROJECT_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
LOGS_DIR = os.path.join(PROJECT_ROOT, 'logs')

# 确保日志目录存在
os.makedirs(LOGS_DIR, exist_ok=True)

# 配置日志
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s',
    handlers=[
        logging.FileHandler(os.path.join(LOGS_DIR, 'ragflow_mcp_tools.log'), mode='a'),
        logging.StreamHandler()
    ]
)

logger = logging.getLogger(__name__)
logger.info("RAGFlow MCP工具模块初始化")

# 加载环境变量
load_dotenv()


# 获取助手列表
def list_assistant() -> str:
    """
    列出所有可用的助手
    :return:
       str: 助手列表信息
    """
    return get_assistant_list()


# 调用大模型返回一个字符串：有可能返回一个助手名称，还有可能返回一个回答的问题内容！
# 第一个参数用户提示词，第二个参数是系统指令
def call_llm(prompt: str, instruction: str) -> str:
    """
    调用大模型进行分析
    :param prompt: 用户提示
    :param instruction: 系统指令
    :return: 大模型的回复
    """
    logger.info(f"调用LLM处理: '{prompt[:50]}...'")
    # 获取开始时间：
    start_time = time.time()
    try:
        # 需要api_key model base_url 等参数
        api_key = os.getenv("DASHSCOPE_API_KEY")
        model = os.getenv("LLM_QWEN2.5")
        base_url = os.getenv("DASHSCOPE_BASE_URL")
        # 需要构建提示词
        prompt_template = ChatPromptTemplate.from_messages([
            ("system", "{instruction}"),
            ("human", "{prompt}")
        ])
        # prompts = "请自我介绍一下你是谁"
        # 获取大模型对象
        llm = ChatOpenAI(api_key=api_key, model=model, base_url=base_url)
        # 创建链：runnable
        llm_chain: Runnable = prompt_template | llm

        # 需要调用链的invoke() 函数
        # 执行链并获取结果
        logger.info(f"开始调用模型 {model}")
        # llm_chain.invoke({}) 不写因为不需要参数，不需要提示词模版
        # 需要参数：具体给系统确定好身份; 用户的提示词是什么
        response = llm_chain.invoke({
            "instruction": instruction,
            "prompt": prompt
        })
        # 将result转换为字符串形式
        result = response.content
        # 获取结束时间
        end_time = time.time()
        # 打印日志
        logger.info(f"成功获取LLM回复，长度: {len(result)}, 耗时: {end_time - start_time:.2f}秒")
        return result

    except Exception as e:
        logger.error(f"LLM调用失败: {str(e)}", exc_info=True)
        return f"调用LLM时发生错误: {str(e)}"


def auto_assistant_answer(question: str) -> str:
    """
    自动选择合适的助手回答问题
    :param question: 用户问题
    :return: 助手的回答
    """
    logger.info(f"自动选择助手回答问题: '{question[:50]}...'")
    # 开始时间
    start_time = time.time()
    try:
        # 1. 获取助手列表
        logger.info("获取助手列表")
        assistants_info = get_assistant_list()
        # 如果字符串为空,或者对应了get_assistant_list的报错返回结果的话,则返回  result = f"获取助手列表时出错: {str(e)}"
        # .startswith()表示字符串是否以指定前缀开头,返回布尔值
        if not assistants_info or assistants_info.startswith("获取助手列表时出错"):
            logger.error("获取助手列表失败")
            return "无法获取助手列表，自动选择助手失败"
        # 2. 构建提示词，让LLM分析选择合适的助手 构建系统提示
        system_prompt = """你是一个助手选择器。你需要根据用户的问题和提供的助手列表，选择最适合回答该问题的助手。

            分析每个助手的功能介绍和关联的知识库，选择最匹配用户问题主题的助手。
            请只返回选择的助手名称，不要添加任何解释或额外文本。"""
        # 用户提示
        user_prompt = f"""用户问题: {question}
        可用的助手列表:
    {assistants_info}
    
    请选择最合适的助手来回答这个问题:"""
        logger.info("调用LLM选择助手")
        # .strip()用于去掉空白字符类型
        # 调用大模型分析具体使用哪个助手
        selected_assistant = call_llm(user_prompt, system_prompt).strip()
        # 打印日志返回的是哪个助手列表
        logger.info(f"助手选择完成，选择的助手: {selected_assistant}")
        # 使用助手助手回答问题
        logger.info(f"使用助手 '{selected_assistant}' 回答问题: '{question[:50]}...'")
        # 调用 create_ask_delete 函数
        answer = create_ask_delete(selected_assistant, question)
        # 获取结束时间
        end_time = time.time()
        logger.info(
            f"自动选择助手回答完成，选择的助手: '{selected_assistant}', 回答长度: {len(answer)}, 总耗时: {end_time - start_time:.2f}秒")
        return answer
    except Exception as e:
        logger.error(f"自动选择助手时出错: {str(e)}", exc_info=True)


# 自动选择助手 assistant_name:助手名称 question:问题 调用RAFGlow 工具中的创建临时会话函数
def choose_assistant(assistant_name: str, question: str) -> str:
    """
    使用指定的助手回答问题
    :param assistant_name: 助手名称
    :param question: 用户问题
    :return: 助手的回答
    """
    logger.info(f"使用指定助手 '{assistant_name}' 回答问题: '{question[:50]}...'")
    # 开始时间
    start_time = time.time()

    try:
        # 调用create_ask_delete并保存返回值
        answer = create_ask_delete(assistant_name, question)
        # 结束时间
        end_time = time.time()
        logger.info(
            f"助手 '{assistant_name}' 回答完成，回答长度: {len(answer) if answer else 0}, 耗时: {end_time - start_time:.2f}秒")
        return answer

    except Exception as e:
        logger.error(f"使用助手 '{assistant_name}' 回答问题时出错: {str(e)}", exc_info=True)
        return f"使用助手回答问题失败: {str(e)}"


if __name__ == "__main__":
    print("RAGFlow MCP工具测试")

    try:
        # 测试列出助手
        # print("\n1. 测试列出助手:")
        # assistants = list_assistant()
        # print(f"助手列表:\n{assistants}\n")

        # # 测试自动选择助手回答问题
        # print("\n2. 测试自动选择助手:")
        # question = "酒驾属于违法行为么?"
        # print(f"问题: {question}")
        # answer = auto_assistant_answer(question)
        # print(f"回答:\n{answer}\n")

        # 测试指定助手回答问题
        print("\n3. 测试指定助手:")
        assistant_name = "刑法知识助手"  # 替换为实际存在的助手名称
        question = "刑法的准则是什么？"
        print(f"助手: {assistant_name}")
        print(f"问题: {question}")
        answer = choose_assistant(assistant_name, question)
        print(f"回答:\n{answer}\n")

    except Exception as e:
        print(f"测试过程中发生错误: {str(e)}")
